{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Object Detection using TAO YOLOv4 with 16-bit imagery\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample prediction of YOLOv4\n",
    "<img align=\"center\" src=\"https://github.com/vpraveen-nv/model_card_images/blob/main/cv/notebook/common/sample.jpg?raw=true\" width=\"960\"> "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Train a ResNet-18 YOLOv4 model on the 16-bit KITTI dataset\n",
    "* Prune the trained YOLOv4 model\n",
    "* Retrain the pruned model to recover lost accuracy\n",
    "* Export the retrained model to .onnx model for inference\n",
    "* Convert the .onnx model to TensorRT engine using tao deploy for inference\n",
    "\n",
    "At the end of this notebook, you will have generated a trained and optimized `yolo_v4` model\n",
    "trained on 16-bit input images, which you may deploy via [Triton](https://github.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps)\n",
    "or [DeepStream](https://developer.nvidia.com/deepstream-sdk).\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of YOLO v4 object detection with 16-bit PNG images using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Install the TAO launcher](#head-1)\n",
    "2. [Prepare dataset and pre-trained model](#head-2) <br>\n",
    "     2.1 [Download the dataset](#head-2-1)<br>\n",
    "     2.2 [Verify the downloaded dataset](#head-2-2)<br>\n",
    "     2.3 [Generate tfrecords](#head-2-3)<br>\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate trained models](#head-5)\n",
    "6. [Prune trained models](#head-6)\n",
    "7. [Retrain pruned models](#head-7)\n",
    "8. [Evaluate retrained model](#head-8)\n",
    "9. [Model Export](#head-9)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users workspace. Please note that the dataset to run this notebook is expected to reside in the `$LOCAL_PROJECT_DIR/data`, while the TAO experiment generated collaterals will be output to `$LOCAL_PROJECT_DIR/yolo_v4`. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.\n",
    "\n",
    "*Note: Please make sure to remove any stray artifacts/files from the `$USER_EXPERIMENT_DIR` or `$DATA_DOWNLOAD_DIR` paths as mentioned below, that may have been generated from previous experiments. Having checkpoint files etc may interfere with creating a training graph for a new experiment.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up env variables for cleaner command line commands.\n",
    "import os\n",
    "\n",
    "%env USER_EXPERIMENT_DIR=/workspace/tao-experiments/yolo_v4_16bit_grayscale\n",
    "%env DATA_DOWNLOAD_DIR=/workspace/tao-experiments/data\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/yolo_v4_16bit_grayscale\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "# The dataset expected to be present in $LOCAL_PROJECT_DIR/data, while the results for the steps\n",
    "# in this notebook will be stored at $LOCAL_PROJECT_DIR/yolo_v4_16bit_grayscale\n",
    "%env LOCAL_PROJECT_DIR=YOUR_LOCAL_PROJECT_DIR_PATH\n",
    "os.environ[\"LOCAL_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\")\n",
    "os.environ[\"LOCAL_EXPERIMENT_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"yolo_v4_16bit_grayscale\")\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"LOCAL_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n",
    "%env SPECS_DIR=/workspace/tao-experiments/yolo_v4_16bit_grayscale/specs\n",
    "\n",
    "# Showing list of specification files.\n",
    "!ls -rlt $LOCAL_SPECS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create local dir\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cell below maps the project directory on your local host to a workspace directory in the TAO docker instance, so that the data and the results are mapped from outside to inside of the docker instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "\n",
    "# Define the dictionary with the mapped drives\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "        # Mapping the data directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "        # Mapping the specs directory.\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "            \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "        },\n",
    "    ]\n",
    "}\n",
    "\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(drive_map, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Install the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n",
    "\n",
    "After setting up your virtual environment with the above requirements, install TAO pip package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install --upgrade nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info --verbose"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset and pre-trained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " We will be using the KITTI detection dataset for the tutorial. To find more details please visit\n",
    " http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=2d. Please download the KITTI detection images (http://www.cvlibs.net/download.php?file=data_object_image_2.zip) and labels (http://www.cvlibs.net/download.php?file=data_object_label_2.zip) to $DATA_DOWNLOAD_DIR.\n",
    " \n",
    " The data will then be extracted to have\n",
    " * training images in `$LOCAL_DATA_DIR/training/image_2`\n",
    " * training labels in `$LOCAL_DATA_DIR/training/label_2`\n",
    " * testing images in `$LOCAL_DATA_DIR/testing/image_2`\n",
    " \n",
    "You may use this notebook with your own dataset as well. To use this example with your own dataset, please follow the same directory structure as mentioned below.\n",
    "\n",
    "*Note: There are no labels for the testing images, therefore we use it just to visualize inferences for the trained model.*"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Download the dataset <a class=\"anchor\" id=\"head-2-1\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once you have gotten the download links in your email, please populate them in place of the `KITTI_IMAGES_DOWNLOAD_URL` and the `KITTI_LABELS_DOWNLOAD_URL`. This next cell, will download the data and place in `$LOCAL_DATA_DIR`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "os.environ[\"URL_IMAGES\"]=KITTI_IMAGES_DOWNLOAD_URL\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_image_2.zip ]; then wget $URL_IMAGES -O $LOCAL_DATA_DIR/data_object_image_2.zip; else echo \"image archive already downloaded\"; fi \n",
    "os.environ[\"URL_LABELS\"]=KITTI_LABELS_DOWNLOAD_URL\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_label_2.zip ]; then wget $URL_LABELS -O $LOCAL_DATA_DIR/data_object_label_2.zip; else echo \"label archive already downloaded\"; fi "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Verify the downloaded dataset <a class=\"anchor\" id=\"head-2-2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the dataset is present\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_image_2.zip ]; then echo 'Image zip file not found, please download.'; else echo 'Found Image zip file.';fi\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_label_2.zip ]; then echo 'Label zip file not found, please download.'; else echo 'Found Labels zip file.';fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This may take a while: verify integrity of zip files \n",
    "!sha256sum $LOCAL_DATA_DIR/data_object_image_2.zip | cut -d ' ' -f 1 | grep -xq '^351c5a2aa0cd9238b50174a3a62b846bc5855da256b82a196431d60ff8d43617$' ; \\\n",
    "if test $? -eq 0; then echo \"images OK\"; else echo \"images corrupt, re-download!\" && rm -f $LOCAL_DATA_DIR/data_object_image_2.zip; fi \n",
    "!sha256sum $LOCAL_DATA_DIR/data_object_label_2.zip | cut -d ' ' -f 1 | grep -xq '^4efc76220d867e1c31bb980bbf8cbc02599f02a9cb4350effa98dbb04aaed880$' ; \\\n",
    "if test $? -eq 0; then echo \"labels OK\"; else echo \"labels corrupt, re-download!\" && rm -f $LOCAL_DATA_DIR/data_object_label_2.zip; fi "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# unpack \n",
    "!unzip -u $LOCAL_DATA_DIR/data_object_image_2.zip -d $LOCAL_DATA_DIR\n",
    "!unzip -u $LOCAL_DATA_DIR/data_object_label_2.zip -d $LOCAL_DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# verify\n",
    "import os\n",
    "\n",
    "DATA_DIR = os.environ.get('LOCAL_DATA_DIR')\n",
    "num_training_images = len(os.listdir(os.path.join(DATA_DIR, \"training/image_2\")))\n",
    "num_training_labels = len(os.listdir(os.path.join(DATA_DIR, \"training/label_2\")))\n",
    "num_testing_images = len(os.listdir(os.path.join(DATA_DIR, \"testing/image_2\")))\n",
    "print(\"Number of images in the train/val set. {}\".format(num_training_images))\n",
    "print(\"Number of labels in the train/val set. {}\".format(num_training_labels))\n",
    "print(\"Number of images in the test set. {}\".format(num_testing_images))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Directory where splitted dataset will be located\n",
    "!mkdir -p $LOCAL_DATA_DIR/kitti_split\n",
    "# Generate val dataset out of training dataset\n",
    "!python3 ../ssd/generate_split.py --input_image_dir=$LOCAL_DATA_DIR/training/image_2 \\\n",
    "                                  --input_label_dir=$LOCAL_DATA_DIR/training/label_2 \\\n",
    "                                  --output_dir=$LOCAL_DATA_DIR/kitti_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert RGB images to (fake) 16-bit grayscale\n",
    "!pip3 install numpy==1.22.2 Pillow==9.0.1\n",
    "import os\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "def to16bit(img_file):\n",
    "    image = Image.open(img_file).convert(\"L\")\n",
    "    # shifted to the higher byte to get a fake 16-bit image\n",
    "    image_np = np.array(image) * 256\n",
    "    image16 = Image.fromarray(image_np.astype(np.uint32))\n",
    "    # overwrite the image file\n",
    "    print(f\"Converting {img_file} to 16-bit grayscale\")\n",
    "    image16.save(img_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate 16-bit grayscale images for train/val splits\n",
    "!mkdir -p $LOCAL_DATA_DIR/kitti_split/training/image_16bit_grayscale\n",
    "!cp $LOCAL_DATA_DIR/kitti_split/training/image/* $LOCAL_DATA_DIR/kitti_split/training/image_16bit_grayscale\n",
    "for img_file in os.listdir(os.path.join(os.environ[\"LOCAL_DATA_DIR\"], \"kitti_split/training/image_16bit_grayscale\")):\n",
    "    image_file = os.path.join(os.environ[\"LOCAL_DATA_DIR\"], \"kitti_split/training/image_16bit_grayscale\", img_file)\n",
    "    to16bit(image_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_DATA_DIR/kitti_split/val/image_16bit_grayscale\n",
    "!cp $LOCAL_DATA_DIR/kitti_split/val/image/* $LOCAL_DATA_DIR/kitti_split/val/image_16bit_grayscale\n",
    "for img_file in os.listdir(os.path.join(os.environ[\"LOCAL_DATA_DIR\"], \"kitti_split/val/image_16bit_grayscale\")):\n",
    "    image_file = os.path.join(os.environ[\"LOCAL_DATA_DIR\"], \"kitti_split/val/image_16bit_grayscale\", img_file)\n",
    "    to16bit(image_file)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally, if you have your own dataset already in a volume (or folder), you can mount the volume on `LOCAL_DATA_DIR` (or create a soft link). Below shows an example:\n",
    "```bash\n",
    "# if your dataset is in /dev/sdc1\n",
    "mount /dev/sdc1 $LOCAL_DATA_DIR\n",
    "\n",
    "# if your dataset is in folder /var/dataset\n",
    "ln -sf /var/dataset $LOCAL_DATA_DIR\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If you use your own dataset, you will need to run the code below to generate the best anchor shape\n",
    "\n",
    "# !tao model yolo_v4 kmeans -l $DATA_DOWNLOAD_DIR/kitti_split/training/label \\\n",
    "#                     -i $DATA_DOWNLOAD_DIR/kitti_split/training/image_16bit_grayscale \\\n",
    "#                     -n 9 \\\n",
    "#                     -x 1248 \\\n",
    "#                     -y 384\n",
    "\n",
    "# The anchor shape generated by this script is sorted. Write the first 3 into small_anchor_shape in the config\n",
    "# file. Write middle 3 into mid_anchor_shape. Write last 3 into big_anchor_shape."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Generate tfrecords <a class=\"anchor\" id=\"head-2-3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model yolo_v4 dataset_convert -d $SPECS_DIR/yolo_v4_tfrecords_kitti_train_16bit_grayscale.txt \\\n",
    "                             -o $DATA_DOWNLOAD_DIR/yolo_v4/tfrecords/train_16bit_grayscale \\\n",
    "                             -r $USER_EXPERIMENT_DIR/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model yolo_v4 dataset_convert -d $SPECS_DIR/yolo_v4_tfrecords_kitti_val_16bit_grayscale.txt \\\n",
    "                             -o $DATA_DOWNLOAD_DIR/yolo_v4/tfrecords/val_16bit_grayscale \\\n",
    "                             -r $USER_EXPERIMENT_DIR/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "* Augmentation parameters for on-the-fly data augmentation\n",
    "* Other training (hyper-)parameters such as batch size, number of epochs, learning rate etc.\n",
    "* Whether to use quantization aware training (QAT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $LOCAL_SPECS_DIR/yolo_v4_train_resnet18_kitti_16bit_grayscale.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the output directory location for models\n",
    "* WARNING: training will take several hours or one day to complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"To run with multigpu, please change --gpus based on the number of available GPUs in your machine.\")\n",
    "!tao model yolo_v4 train -e $SPECS_DIR/yolo_v4_train_resnet18_kitti_16bit_grayscale.txt \\\n",
    "                   -r $USER_EXPERIMENT_DIR/experiment_dir_unpruned \\\n",
    "                   --gpus 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Model for each epoch:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned/weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now check the evaluation stats in the csv file and pick the model with highest eval accuracy.\n",
    "!cat $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned/yolov4_training_log_resnet18.csv\n",
    "%set_env EPOCH=080"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate trained models <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model yolo_v4 evaluate -e $SPECS_DIR/yolo_v4_train_resnet18_kitti_16bit_grayscale.txt \\\n",
    "                      -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/weights/yolov4_resnet18_epoch_$EPOCH.hdf5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Prune trained models <a class=\"anchor\" id=\"head-6\"></a>\n",
    "* Specify pre-trained model\n",
    "* Equalization criterion (`Only for resnets as they have element wise operations or MobileNets.`)\n",
    "* Threshold for pruning.\n",
    "* A key to save and load the model\n",
    "* Output directory to store the model\n",
    "\n",
    "Usually, you just need to adjust `-pth` (threshold) for accuracy and model size trade off. Higher `pth` gives you smaller model (and thus higher inference speed) but worse accuracy. The threshold value depends on the dataset and the model. `0.5` in the block below is just a start point. If the retrain accuracy is good, you can increase this value to get smaller models. Otherwise, lower this value to get better accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_pruned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model yolo_v4 prune -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/weights/yolov4_resnet18_epoch_$EPOCH.hdf5 \\\n",
    "                   -e $SPECS_DIR/yolo_v4_train_resnet18_kitti_16bit_grayscale.txt \\\n",
    "                   -o $USER_EXPERIMENT_DIR/experiment_dir_pruned/yolov4_resnet18_pruned.hdf5 \\\n",
    "                   -eq intersection \\\n",
    "                   -pth 0.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/experiment_dir_pruned/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Retrain pruned models <a class=\"anchor\" id=\"head-7\"></a>\n",
    "* Model needs to be re-trained to bring back accuracy after pruning\n",
    "* Specify re-training specification\n",
    "* WARNING: training will take several hours or one day to complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Printing the retrain spec file. \n",
    "# Here we have updated the spec file to include the newly pruned model as a pretrained weights.\n",
    "!sed -i 's,EXPERIMENT_DIR,'\"$USER_EXPERIMENT_DIR\"',' $LOCAL_SPECS_DIR/yolo_v4_retrain_resnet18_kitti_16bit_grayscale.txt\n",
    "!cat $LOCAL_SPECS_DIR/yolo_v4_retrain_resnet18_kitti_16bit_grayscale.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retraining using the pruned model as pretrained weights\n",
    "!tao model yolo_v4 train --gpus 1 \\\n",
    "                   -e $SPECS_DIR/yolo_v4_retrain_resnet18_kitti_16bit_grayscale.txt \\\n",
    "                   -r $USER_EXPERIMENT_DIR/experiment_dir_retrain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Listing the newly retrained model.\n",
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain/weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now check the evaluation stats in the csv file and pick the model with highest eval accuracy.\n",
    "!cat $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain/yolov4_training_log_resnet18.csv\n",
    "%set_env EPOCH=080"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Evaluate retrained model <a class=\"anchor\" id=\"head-8\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model yolo_v4 evaluate -e $SPECS_DIR/yolo_v4_retrain_resnet18_kitti_16bit_grayscale.txt \\\n",
    "                      -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/yolov4_resnet18_epoch_$EPOCH.hdf5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Model Export <a class=\"anchor\" id=\"head-9\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you trained a non-QAT model, you may export in FP32, FP16 or INT8 mode using the code block below. For INT8, you need to provide calibration image directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# tao <task> export will fail if .onnx already exists. So we clear the export folder before tao <task> export\n",
    "!rm -rf $LOCAL_EXPERIMENT_DIR/export\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/export\n",
    "# Generate .onnx file using tao container\n",
    "!tao model yolo_v4 export -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/yolov4_resnet18_epoch_$EPOCH.hdf5 \\\n",
    "                    -o $USER_EXPERIMENT_DIR/export/yolov4_resnet18_epoch_$EPOCH.onnx \\\n",
    "                    -e $SPECS_DIR/yolo_v4_retrain_resnet18_kitti_16bit_grayscale.txt \\\n",
    "                    --target_opset 12 \\\n",
    "                    --gen_ds_config"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the `tao deploy` container, you can generate a TensorRT engine and verify the correctness of the generated through evaluate and inference.\n",
    "\n",
    "The `tao deploy` produces optimized tensorrt engines for the platform that it resides on. Therefore, to get maximum performance, please run `tao deploy` command which will instantiate a deploy container, with the exported `.onnx` file on your target device. The `tao deploy` container only works for x86, with discrete NVIDIA GPU's.\n",
    "\n",
    "For the jetson devices, please download the tao-converter for jetson and refer to [here](https://docs.nvidia.com/tao/tao-toolkit-archive/tao-30-2108/text/tensorrt.html#installing-the-tao-converter) for more details.\n",
    "\n",
    "If you choose to integrate your model into deepstream directly, you may do so by simply copying the exported `.onnx` file along with the calibration cache to the target device and updating the spec file that configures the `gst-nvinfer` element to point to this newly exported model. Usually this file is called `config_infer_primary.txt` for detection models and `config_infer_secondary_*.txt` for classification models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert to TensorRT engine (FP32). Change --data_type to fp16 for FP16 mode\n",
    "!tao deploy yolo_v4 gen_trt_engine -m $USER_EXPERIMENT_DIR/export/yolov4_resnet18_epoch_$EPOCH.onnx \\\n",
    "                                   -e $SPECS_DIR/yolo_v4_retrain_resnet18_kitti_16bit_grayscale.txt \\\n",
    "                                   --batch_size 16 \\\n",
    "                                   --min_batch_size 1 \\\n",
    "                                   --opt_batch_size 8 \\\n",
    "                                   --max_batch_size 16 \\\n",
    "                                   --data_type fp32 \\\n",
    "                                   --engine_file $USER_EXPERIMENT_DIR/export/trt.engine \\\n",
    "                                   --results_dir $USER_EXPERIMENT_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert to TensorRT engine (INT8).\n",
    "!tao deploy yolo_v4 gen_trt_engine -m $USER_EXPERIMENT_DIR/export/yolov4_resnet18_epoch_$EPOCH.onnx \\\n",
    "                                   -e $SPECS_DIR/yolo_v4_retrain_resnet18_kitti_16bit_grayscale.txt \\\n",
    "                                   --cal_image_dir $DATA_DOWNLOAD_DIR/kitti_split/training/image_16bit_grayscale \\\n",
    "                                   --data_type int8 \\\n",
    "                                   --batch_size 16 \\\n",
    "                                   --min_batch_size 1 \\\n",
    "                                   --opt_batch_size 8 \\\n",
    "                                   --max_batch_size 16 \\\n",
    "                                   --batches 10 \\\n",
    "                                   --cal_cache_file $USER_EXPERIMENT_DIR/export/cal.bin  \\\n",
    "                                   --cal_data_file $USER_EXPERIMENT_DIR/export/cal.tensorfile \\\n",
    "                                   --engine_file $USER_EXPERIMENT_DIR/export/trt.engine \\\n",
    "                                   --results_dir $USER_EXPERIMENT_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model:')\n",
    "print('------------')\n",
    "!ls -lh $LOCAL_EXPERIMENT_DIR/export"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
